import os
import numpy as np
import cv2
import re


def load_flow(flow_dir, row, col):
    filename = f"tr{row}-tc{col}.npy"
    filepath = os.path.join(flow_dir, filename)
    return np.load(filepath, mmap_mode='r')


def load_image(image_dir, row, col):
    # 正则表达式匹配前缀部分和tr-row-tc-col部分
    pattern = re.compile(r'(.*)_tr(\d+)-tc(\d+)\.png')

    # 遍历文件夹中的所有文件
    for filename in os.listdir(image_dir):
        # 检查文件名是否符合完整的tr{row}-tc{col}格式
        match = pattern.match(filename)
        if match:
            prefix = match.group(1)
            matched_row = int(match.group(2))
            matched_col = int(match.group(3))
            if matched_row == row and matched_col == col:
                filepath = os.path.join(image_dir, filename)
                image = cv2.imread(filepath)
                if image is not None:
                    return image, prefix, filename
                else:
                    print(f"Error reading image at {filepath}")

    # 如果没有匹配的文件，返回None
    return None, None, None

def save_flow(flow, row, col, output_dir):
    filename = f"tr{row}-tc{col}.npy"
    filepath = os.path.join(output_dir, filename)
    np.save(filepath, flow)
    print(f"Saved {filename} to {filepath}")


def save_image(image, filename, output_dir):
    filepath = os.path.join(output_dir, filename)
    cv2.imwrite(filepath, image)
    print(f"Saved {filename} to {filepath}")


def combine_flows_and_images(flow_dir, image_dir, flow_output_dir, image_output_dir):
    rows, cols = 37, 33
    flow_tile_size = 1024
    image_tile_size = 1024
    flow_subgrid_rows, flow_subgrid_cols = 15, 13
    row_overlap, col_overlap = 4, 3

    os.makedirs(flow_output_dir, exist_ok=True)
    os.makedirs(image_output_dir, exist_ok=True)

    prefix = None
    processed_subgrids = set()

    for subgrid_row in range(1, rows + 1, flow_subgrid_rows - row_overlap):
        if subgrid_row + flow_subgrid_rows - 1 > rows:
            subgrid_row = rows - flow_subgrid_rows + 1
        for subgrid_col in range(1, cols + 1, flow_subgrid_cols - col_overlap):
            if subgrid_col + flow_subgrid_cols - 1 > cols:
                subgrid_col = cols - flow_subgrid_cols + 1

            subgrid_id = (subgrid_row, subgrid_col)
            if subgrid_id in processed_subgrids:
                continue
            processed_subgrids.add(subgrid_id)

            flow_combined = np.zeros((flow_subgrid_rows * flow_tile_size, flow_subgrid_cols * flow_tile_size, 2),
                                     dtype=np.float32)
            image_combined = np.zeros((flow_subgrid_rows * image_tile_size, flow_subgrid_cols * image_tile_size, 3),
                                      dtype=np.uint8)

            for i in range(subgrid_row, subgrid_row + flow_subgrid_rows):
                for j in range(subgrid_col, subgrid_col + flow_subgrid_cols):
                    flow_part = load_flow(flow_dir, i, j)
                    row_offset = (i - subgrid_row) * flow_tile_size
                    col_offset = (j - subgrid_col) * flow_tile_size
                    flow_combined[row_offset:row_offset + flow_tile_size, col_offset:col_offset + flow_tile_size,
                    :] = flow_part

                    image_part, temp_prefix, filename = load_image(image_dir, i, j)
                    if prefix is None:
                        prefix = temp_prefix
                    image_combined[row_offset:row_offset + image_tile_size, col_offset:col_offset + image_tile_size,
                    :] = image_part

            save_flow(flow_combined, (subgrid_row - 1) // (flow_subgrid_rows - row_overlap) + 1,
                      (subgrid_col - 1) // (flow_subgrid_cols - col_overlap) + 1, flow_output_dir)
            save_image(image_combined,
                       f"{prefix}_tr{(subgrid_row - 1) // (flow_subgrid_rows - row_overlap) + 1}-tc{(subgrid_col - 1) // (flow_subgrid_cols - col_overlap) + 1}.png",
                       image_output_dir)


# 示例用法
flow_dir = 'E:\\wafer52\\32nm_flow_fine_align(3)'  # 替换为实际光流场文件夹路径
image_dir = 'E:\\wafer52\\11867_32nm_coarse_align2(1)'  # 替换为实际图片文件夹路径
flow_output_dir = 'E:\\wafer52\\11867_32nm_fine_flow_big_combine(3)'  # 替换为实际光流场合并输出文件夹路径
image_output_dir = 'E:\\wafer52\\11867_32nm_big_coarse_align2_combine(1)'  # 替换为实际图片合并输出文件夹路径


combine_flows_and_images(flow_dir, image_dir, flow_output_dir, image_output_dir)
